import math
import os
import sys
import torch
import time
from dataprocess.datareprepare import prepare_data_loaders
from model.dborgotmodel import dborgotmodel
from model.dbmodel import dbmodel
from model.gotmodel import gotmdmodel
from utils.trainresultdraw import plot_accuracy_and_loss
from torch import nn
from torch.utils.data import TensorDataset, DataLoader, random_split
from utils.fldbtrainprocess import fldbtrain
from utils.dborgottrainprocess import dborgottrain

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = dborgotmodel().to(device)
dbmodel = dbmodel().to(device)
gotmdmodel = gotmdmodel().to(device)
Loss_fn = nn.CrossEntropyLoss().to(device)
Yh_fn = torch.optim.SGD(model.parameters(), lr=5e-4, momentum=0.9)
db_Yh_fn = torch.optim.SGD(dbmodel.parameters(), lr=5e-4, momentum=0.9)
gotmd_Yh_fn = torch.optim.SGD(gotmdmodel.parameters(), lr=5e-4, momentum=0.9)
train_loader, test_loader = prepare_data_loaders()
t_a, t_loss, all_x_db, all_y_db, all_x_gotmd, all_y_gotmd = dborgottrain(train_loader, model, Loss_fn, Yh_fn)

print('疾病分类预测正确率:' + str(t_a) + '，训练Loss:' + str(t_loss))
all_x_db_tensor = torch.cat(all_x_db, dim=0)
all_y_db_tensor = torch.cat(all_y_db, dim=0)
all_x_gotmd_tensor = torch.cat(all_x_gotmd, dim=0)
all_y_gotmd_tensor = torch.cat(all_y_gotmd, dim=0)
x_db_length = all_x_db_tensor.size(0)
x_gotmd_length = all_x_gotmd_tensor.size(0)
nearest_multiple_128 = int(128 * math.floor(x_db_length / 128.0))
gotmd_nearest_multiple_128 = int(128 * math.floor(x_gotmd_length / 128.0))
all_x_db_tensor = all_x_db_tensor[:nearest_multiple_128]
all_y_db_tensor = all_y_db_tensor[:nearest_multiple_128]
all_x_gotmd_tensor = all_x_db_tensor[:gotmd_nearest_multiple_128]
all_y_gotmd_tensor = all_y_db_tensor[:gotmd_nearest_multiple_128]
gotmd_dataset = TensorDataset(all_x_gotmd_tensor, all_y_gotmd_tensor)
db_dataset = TensorDataset(all_x_db_tensor, all_y_db_tensor)
db_dataset = DataLoader(dataset=db_dataset, batch_size=128, shuffle=True)
gotmd_dataset = DataLoader(dataset=gotmd_dataset, batch_size=128, shuffle=True)
# 划分为两个客户端
# print(len(db_dataset))
# print(len(gotmd_dataset))
# print(len(train_loader))
# exit()
n = len(db_dataset)
split_lengths = [n // 2, n - n // 2]
client_datasets = random_split(db_dataset, split_lengths)
client_loaders = [DataLoader(client_dataset, batch_size=8) for client_dataset in client_datasets]


# print(len(client_loaders[0]), len(client_loaders[1]), len(train_loader))
# exit()
# 设定联邦学习轮数
def communication_overhead(client_params_list):
    total_size = 0
    for params in client_params_list:
        for key in params.keys():
            tensor_size = sys.getsizeof(params[key])
            total_size += tensor_size
    return total_size


num_rounds = 10
# 用于存储每个轮次的训练时间
round_training_times = []
for round in range(num_rounds):
    round_start_time = time.time()
    print(f"Round {round + 1} of federated learning")
    client_params_list = []
    client_training_times = []
    # 客户端进行本地训练并返回参数
    for index, client_loader in enumerate(client_loaders):
        client_start_time = time.time()
        modelState, acc, loss = fldbtrain(db_dataset, dbmodel, Loss_fn, db_Yh_fn, round)
        client_end_time = time.time()
        client_training_time = client_end_time - client_start_time
        client_training_times.append(client_training_time)
        print(f"客户端{index + 1}第{round + 1}轮的训练时间为: {client_training_time}秒")
        print(f"客户端{index + 1}第{round + 1}轮的训练准确率{acc}损失{loss}")
        client_params_list.append(modelState)
    aggregated_params = {}
    for key in client_params_list[0].keys():
        aggregated_params[key] = torch.stack([params[key] for params in client_params_list]).mean(dim=0)
    params = communication_overhead(client_params_list)
    print(f"参数规模{params}")
    dbmodel.load_state_dict(aggregated_params)
    round_end_time = time.time()
    round_training_time = round_end_time - round_start_time
    round_training_times.append(round_training_time)
    print(f"第{round + 1}轮联邦学习总训练时间为: {round_training_time}秒")
average_round_training_time = sum(round_training_times)/len(round_training_times)
print(f"平均每轮训练时间为: {average_round_training_time}秒")